Watershed-based tree segmentation¶

Watershed segmentation, also known as watershed transformation, is an image processing technique used to segment regions of interest in digital images. This technique is widely used in computer vision and image processing for object detection and segmentation, especially in images where the boundaries of objects are not well defined.

The main idea of ​​watershed segmentation is to treat the image as a topographic map, where peaks represent regions that need to be segmented. The process involves analyzing intensity features of the image to identify regions of interest and separating them into different regions or segments.

These are the main steps of the watershed segmentation algorithm:

  • Gradient transformation: First, the gradient of the image is calculated. This is done to highlight the edges or boundaries between objects in the image. There are several ways to calculate the gradient, such as the Sobel or Canny operator.

  • Marker identification: Next, initial markers are chosen, which correspond to well-defined points or areas of interest in the image. These markers can be chosen manually or automatically depending on the application.

  • Watershed transformation: The watershed transformation is applied to the identified markers and the gradient of the image. In this process, the markers act as seeds that generate a flood in the image. The flood starts at the markers and spreads throughout the image, creating separate regions and associating pixels with those regions.

  • Final segmentation: After the watershed transformation we obtain an initial segmentation of the image. However, this segmentation may contain overlapping regions. To correct this, a region merging step is performed based on specific criteria such as distance or intensity difference.

Watersheed segmentation is particularly useful in images with overlapping objects or in regions where the edges are poorly defined. However, one of the challenges of this technique is its sensitivity to noise, which can lead to inaccurate segmentations. To overcome this problem, pre- and post-processing can be applied to improve the quality of the results.

In addition, the watershed transform can also be combined with other segmentation techniques, such as region-based segmentation or region-growing segmentation, to obtain better results in different types of images and applications.

Let's apply watersheed segmentation to a Drone image to obtain the polygons of each plant.

In [ ]:
!pip install rasterio

Collecting rasterio
  Downloading rasterio-1.3.10-cp310-cp310-manylinux2014_x86_64.whl.metadata (14 kB)
Collecting affine (from rasterio)
  Downloading affine-2.4.0-py3-none-any.whl.metadata (4.0 kB)
Requirement already satisfied: attrs in /usr/local/lib/python3.10/dist-packages (from rasterio) (24.2.0)
Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from rasterio) (2024.7.4)
Requirement already satisfied: click>=4.0 in /usr/local/lib/python3.10/dist-packages (from rasterio) (8.1.7)
Requirement already satisfied: cligj>=0.5 in /usr/local/lib/python3.10/dist-packages (from rasterio) (0.7.2)
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from rasterio) (1.26.4)
Collecting snuggs>=1.4.1 (from rasterio)
  Downloading snuggs-1.4.7-py3-none-any.whl.metadata (3.4 kB)
Requirement already satisfied: click-plugins in /usr/local/lib/python3.10/dist-packages (from rasterio) (1.1.1)
Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from rasterio) (71.0.4)
Requirement already satisfied: pyparsing>=2.1.6 in /usr/local/lib/python3.10/dist-packages (from snuggs>=1.4.1->rasterio) (3.1.2)
Downloading rasterio-1.3.10-cp310-cp310-manylinux2014_x86_64.whl (21.5 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 21.5/21.5 MB 37.3 MB/s eta 0:00:00
Downloading snuggs-1.4.7-py3-none-any.whl (5.4 kB)
Downloading affine-2.4.0-py3-none-any.whl (15 kB)
Installing collected packages: snuggs, affine, rasterio
Successfully installed affine-2.4.0 rasterio-1.3.10 snuggs-1.4.7
In [ ]:
from google.colab import drive
drive.mount('/content/drive')

Mounted at /content/drive
In [ ]:
from skimage import io
from skimage.feature import peak_local_max
from scipy import ndimage as ndi
from skimage.segmentation import watershed
import matplotlib.pyplot as plt
import numpy as np
import cv2
import rasterio
from skimage.segmentation import mark_boundaries
from rasterio.plot import show

Let's set the image path, open and display it with matplotlib:

In [ ]:
path_img = '/content/drive/MyDrive/Datasets_CV/AOI_teste.tif'
In [ ]:
src = rasterio.open(path_img)
In [ ]:
img = src.read()
In [ ]:
img = img.transpose([1,2,0])
In [ ]:
img = img.astype('uint8')
In [ ]:
plt.figure(figsize=[16,16])
plt.imshow(img)
Out[ ]:
<matplotlib.image.AxesImage at 0x788495fb1510>
No description has been provided for this image

We can convert the image to BGR with opencv and apply a median blur filter to this image:

In [ ]:
bgr_img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
In [ ]:
resulting_image = cv2.medianBlur(bgr_img ,35)
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(resulting_image)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image

We will then convert this image to HSV color space and filter out the green color. We will create a mask with the green color.

In [ ]:
hsv_img = cv2.cvtColor(resulting_image, cv2.COLOR_BGR2HSV)
In [ ]:
mask = cv2.inRange(hsv_img, (36, 50, 50), (70, 255,255))
imask = mask>0
green = np.zeros_like(img, np.uint8)
green[imask] = img[imask]
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(green)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image

With the mask we will apply morphological filters to generate regions in the image:

In [ ]:
mask = np.where(mask == 255, 1,0)
In [ ]:
mask = mask.astype('uint8')
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(mask)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image
In [ ]:
dilate_se=cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
result = cv2.dilate(mask,dilate_se,iterations = 5)
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(result)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image
In [ ]:
erosion_se=cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
erosion = cv2.erode(result,erosion_se,iterations = 28)
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(erosion)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image

We will apply the distance transform to obtain the image of the regions with the distance from the edge:

In [ ]:
dist_transform = cv2.distanceTransform(erosion, cv2.DIST_L2,3)
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(dist_transform)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image

Then we filter the core areas:

In [ ]:
ret, sure_fg = cv2.threshold(dist_transform,0.3*dist_transform.max(),255,0)
In [ ]:
plt.figure(figsize=[12,12])
plt.imshow(sure_fg)
plt.axis('off')
Out[ ]:
(-0.5, 4555.5, 6263.5, -0.5)
No description has been provided for this image

We obtain the unique segments of these core areas and apply the Watershed algorithm:

In [ ]:
markers, _ = ndi.label(sure_fg)
In [ ]:
segmented = watershed(255-dist_transform, markers, mask=mask)

With the generated segments we convert them into polygons using the georeferencing of the original image as a base:

In [ ]:
segments_quick = segmented.astype('uint16')
In [ ]:
segments_quick
Out[ ]:
array([[0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       ...,
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0]], dtype=uint16)
In [ ]:
mask_seg = np.where(segments_quick == 0, 0,1)
mask_seg = mask_seg.astype('uint8')
In [ ]:
from rasterio import features
from shapely.geometry import shape
import geopandas as gpd
import shapely
In [ ]:
find_shapes = features.shapes(segments_quick, mask=mask_seg,  transform=src.transform)
In [ ]:
polygons = [shapely.geometry.Polygon(shape[0]["coordinates"][0]) for shape in find_shapes]
In [ ]:
crs = src.crs
Poly_gdf = gpd.GeoDataFrame(crs=crs, geometry=polygons)
In [ ]:
Poly_gdf
Out[ ]:
geometry
0 POLYGON ((394372.805 7306534.973, 394372.805 7...
1 POLYGON ((394377.860 7306528.161, 394377.860 7...
2 POLYGON ((394374.110 7306531.056, 394374.110 7...
3 POLYGON ((394369.410 7306527.828, 394369.410 7...
4 POLYGON ((394377.908 7306528.256, 394377.908 7...
... ...
238 POLYGON ((394439.217 7306410.384, 394439.217 7...
239 POLYGON ((394430.791 7306408.889, 394430.791 7...
240 POLYGON ((394442.754 7306406.895, 394442.754 7...
241 POLYGON ((394433.710 7306405.423, 394433.710 7...
242 POLYGON ((394436.986 7306402.931, 394436.986 7...

243 rows × 1 columns

We present the polygons:

In [ ]:
Poly_gdf.boundary.plot(figsize=(20,14))
Out[ ]:
<Axes: >
No description has been provided for this image

We can present the polygons together with the image:

In [ ]:
Poly_gdf.reset_index(inplace=True)
In [ ]:
fig, ax = plt.subplots(figsize=(15, 15))
with rasterio.open(path_img) as src:
    show(src,ax=ax)
Poly_gdf.plot('index', ax=ax, cmap='hsv', alpha=0.5)
Poly_gdf.boundary.plot(ax=ax, edgecolor='red')
Out[ ]:
<Axes: >
No description has been provided for this image